from __future__ import print_function


from pylab import *
import sys
import StringIO

from camori import EdgelSet, quaternion_to_matrix, quaternion_product_c as qp
import simplejson
import scipy.io

from quaternion import *


def matrix_to_quaternion(Q):
    T = 1. + trace(Q)

    if T > 1e-8:
        W = sqrt(T/4.)
        X = (Q[1,2]-Q[2,1]) / (4. * W)
        Y = (Q[2,0]-Q[0,2]) / (4. * W)
        Z = (Q[0,1]-Q[1,0]) / (4. * W)

        # print(W*W+X*X+Y*Y+Z*Z)
        nf = (W*W+X*X+Y*Y+Z*Z)**-.5
        return array((X,Y,Z))*nf

    # else
    perm = array([[0,1,2], [1,2,0], [2,0,1]])
    p = perm[argmax(diag(Q))]

    u,v,w = p

    xx = 1 + Q[u,u] - Q[v,v] - Q[w,w]

    if xx <= 0:
        return zeros(3)

    r = np.sqrt(xx)
    
    q0 = (Q[w,v] - Q[v,w]) / (2 * r)
    qu = r / 2
    qv = (Q[u,v] + Q[v,u]) / (2 * r)
    qw = (Q[w,u] + Q[u,w]) / (2 * r)

    nf = (q0*q0+qu*qu+qv*qv+qw*qw)**-0.5


    out = zeros(3)
    out[p] = array([qu,qv,qw])*nf
    
    return out

if __name__ == '__main__':
    ion()

    ## Reads the first 102 lines of the text, one for each of the 102
    ## images in the yorkurbandb dataset.
    ff = ''.join(open(sys.argv[1]).readlines()[:102])
    in_txt = StringIO.StringIO(ff)
    myqL = np.loadtxt(in_txt)

    ## Sort input
    myqL = myqL[argsort(myqL[:,0])]


    set_printoptions(precision=3)

    names = (x[:-1] for x in open('yorknames.txt').readlines())

    Rs = [array([[ 1, 0, 0],
                 [ 0, 1, 0],
                 [ 0, 0, 1] ]),
          array([[ 1, 0, 0],
                 [ 0, 0, 1],
                   [ 0,-1, 0] ]),
          array([[ 0, 0, 1],
                 [ 0, 1, 0],
                 [-1, 0, 0] ]),
          array([[ 0, 1, 0],
                 [-1, 0, 0],
                 [ 0, 0, 1] ]),
          ]

    rr = set([tuple(dot(dot(x,y),z).ravel()) for x in Rs  for y in Rs  for z in Rs])
    rr = list(rr)

    rr.sort()
    rr.reverse()

    rr = [reshape(array(r),(3,3)) for r in rr]

    def normalize(M):
        L = [matrix_to_quaternion(dot(M,r)) for r in rr]
        return L
    
    root = '/home/nlw/ciencia/DADOS/dists/YorkUrbanDB/'
    name = sys.argv[1]

    allerrs1=[]
    allerrs2=[]
    alltt=[]


    for frm,who in enumerate(names):
        ## Red output from edgel_orientation
        idx = myqL[frm,0]
        myq1 = Quat(myqL[frm,1:5])
        myq2 = Quat(myqL[frm,5:9])
        tt = myqL[frm,9:]

        ## Get camera orientation (vps = "vanishing points")
        gt_vp=scipy.io.loadmat('/home/nlw/ciencia/DADOS/dists/YorkUrbanDB/%s/%sGroundTruthVP_Orthogonal_CamParams.mat'%(who,who))
        vps=gt_vp['vp_orthogonal']
    ## Convert for downwards y axis
        vps = dot(dot(array([[1,0,0],[0,-1,0],[0,0,1]]) , vps), array([[1,0,0],[0,-1,0],[0,0,1]]))

        q = Quat(matrix_to_quaternion(vps.T))

        err1 = (myq1/q).canonical().angle()
        err2 = (myq2/q).canonical().angle()

        print(frm, err1, err2)
        print('ref:', q.canonical())
        print('est1:', myq1.canonical())
        print('est2:', myq2.canonical())
        print()
        allerrs1.append(err1)
        allerrs2.append(err2)
        alltt.append(tt)




    allerrs1 = array(allerrs1)
    allerrs2 = array(allerrs2)
    alltt = array(alltt)

    figure(1)
    sallerrs1 = sort(allerrs1)
    sallerrs2 = sort(allerrs2)
    p = mgrid[:sallerrs1.shape[0]] * 100.  / (sallerrs1.shape[0]-1)
    plot(sallerrs1, p, '-+' )
    plot(sallerrs2, p, '-+' )
    text(sallerrs2[-1], 98,  nonzero(allerrs2 == sallerrs2[-1])[0][0],va='top',ha='center')
    text(sallerrs2[-2], 98,  nonzero(allerrs2 == sallerrs2[-2])[0][0],va='top',ha='center')
    text(sallerrs2[-3], 95,  nonzero(allerrs2 == sallerrs2[-3])[0][0],va='top',ha='center')
    text(sallerrs2[-4], 92,  nonzero(allerrs2 == sallerrs2[-4])[0][0],va='top',ha='center')
    axis([0,10,0,102])
    grid()

    ylabel('Imagem')
    xlabel('Erro [graus]')


    print(70*'=')
    print('minerr', sallerrs1[0])
    # print('10therr', sallerrs1[9])
    # print('50therr', sallerrs1[49])
    # print('90therr', sallerrs1[89])
    print('1Qerr', sallerrs1[25])
    print('miderr', sallerrs1[50])
    print('3Qerr', sallerrs1[76])
    print('maxerr', sallerrs1[-1])
    print('errmean', mean(allerrs1))
    print('errstd', std(allerrs1))

    print(70*'=')
    print('minerr', sallerrs2[0])
    # print('10therr', sallerrs2[9])
    # print('50therr', sallerrs2[49])
    # print('90therr', sallerrs2[89])
    print('1Qerr', '%.2f'%sallerrs2[25])
    print('miderr', '%.2f'%sallerrs2[50])
    print('3Qerr', '%.2f'%sallerrs2[76])
    print('maxerr', sallerrs2[-1])
    print('errmean', '%.2f'%mean(allerrs2))
    print('errstd', '%.2f'%std(allerrs2))

    print(70*'=')
    print('Ned_min', alltt[:,0].min())
    print('Ned_mean', alltt[:,0].mean())
    print('Ned_max', alltt[:,0].max())

    print(70*'=')
    print('rand_it_mean', mean(alltt[:,1]))
    print('rand_std', std(alltt[:,1]))
    print('fsqp_it_mean', mean(alltt[:,2]))
    print('fsqp_std', std(alltt[:,2]))

    print(70*'=')
    print('rand_time_mean', mean(alltt[:,5]))
    print('rand_time_std', std(alltt[:,5]))
    print('fsqp_time_mean', mean(alltt[:,6]))
    print('fsqp_time_std', std(alltt[:,6]))
    print('topt_time_mean', '%.2f'%mean(alltt[:,5]+alltt[:,6]))
    print('topt_time_std', std(alltt[:,5]+alltt[:,6]))


    print(mean(alltt[:,5]), std(alltt[:,5]), mean(alltt[:,5]+alltt[:,6]), std(alltt[:,5]+alltt[:,6]))

    # print(sallerrs2[0],sallerrs2[25],sallerrs2[50],sallerrs2[76],sallerrs2[-1],mean(allerrs2),mean(alltt[:,5]+alltt[:,6]),std(alltt[:,5]+alltt[:,6]))
